import numpy as np
from sklearn import datasets, metrics, svm
from sklearn.model_selection import train_test_split
from skimage import io, color, transform
import matplotlib.pyplot as plt

# 加载 MNIST 数据集
mnist = datasets.load_digits()

# 数据集预处理
images = mnist.images
labels = mnist.target

# 展平图像
images = images.reshape(len(images), -1)

# 将数据集划分为 3:1 的训练集和测试集
X_train, X_test, y_train, y_test = train_test_split(images, labels, test_size=0.25, random_state=42)

# 初始化支持向量机分类器
SVM_classifier = svm.SVC(gamma=0.001, probability=True)

# 训练模型
SVM_classifier.fit(X_train, y_train)

# 在测试集上进行预测
y_pred = SVM_classifier.predict(X_test)

# 显示分类报告
print("Performance Report:\n", metrics.classification_report(y_test, y_pred))


# 定义函数，加载自定义图片并预测
def predict_custom_image(image_path):
    # 加载图片
    digit_img = io.imread(image_path)
    # 转为灰度图
    digit_img = color.rgb2gray(color.rgba2rgb(digit_img))
    # 调整大小为 8x8
    digit_img = transform.resize(digit_img, (8, 8), mode="reflect")
    # 归一化为 0-16 之间的像素值，与原始数据集匹配
    digit_img = digit_img * 16
    # 展平为 1D 数组
    digit_flattened = digit_img.flatten().reshape(1, -1)

    # 预测图片
    prediction = SVM_classifier.predict(digit_flattened)
    probabilities = SVM_classifier.predict_proba(digit_flattened)

    print(f"Predicted class: {prediction[0]}")
    print(f"Prediction probabilities: {probabilities}")


# 测试自定义图片
predict_custom_image('images/digit_6.png')
